import os
import cv2
import numpy as np
from multiprocessing.pool import  Pool

target_path = '/my_data'

surf_detector = cv2.xfeatures2d.SURF_create()
std_height = 1024
matcher = cv2.BFMatcher(cv2.NORM_L1, crossCheck=True)
def aling(img_name):
    im_path = os.path.join(img_path, img_name)
    im = cv2.cvtColor(cv2.imread(im_path), cv2.COLOR_BGR2GRAY)

    im_path = os.path.join(ori_temp_path, img_name.replace('.jpg', '_t.jpg'))
    rgb_temp = cv2.imread(im_path)
    im_t = cv2.cvtColor(rgb_temp, cv2.COLOR_BGR2GRAY)

    im_t = (im_t + (np.mean(im) - np.mean(im_t))).astype(np.uint8)

    height, width = im_t.shape[:2]
    h_ratio = 1.0 * height / std_height
    std_width = int(round(width / h_ratio))
    w_ratio = 1.0 * width / std_width

    kp1, d1 = surf_detector.detectAndCompute(cv2.resize(im, (std_width, std_height)), None)
    kp2, d2 = surf_detector.detectAndCompute(cv2.resize(im_t, (std_width, std_height)), None)

    matches = matcher.match(d1, d2)
    matches = [match for match in matches if match.distance <= 0.75]
    no_of_matches = len(matches)

    print("num of matches", no_of_matches)
    if no_of_matches >= 10:
        # Define empty matrices of shape no_of_matches * 2.
        p1 = np.zeros((no_of_matches, 2))
        p2 = np.zeros((no_of_matches, 2))

        for i in range(len(matches)):
            p1[i, :] = kp1[matches[i].queryIdx].pt
            p2[i, :] = kp2[matches[i].trainIdx].pt

        p1 = p1 * [[w_ratio, h_ratio]]
        p2 = p2 * [[w_ratio, h_ratio]]

        homography, mask = cv2.findHomography(p2, p1, cv2.RANSAC)
        transformed_img = cv2.warpPerspective(rgb_temp, homography, (width, height))
    else:
        transformed_img = rgb_temp
    cv2.imwrite(os.path.join(fast_align_temp_path, img_name.replace('.jpg', '_t.jpg')), transformed_img)

if __name__ == '__main__':
    img_path = '/my_data/ori_img'
    ori_temp_path = '/my_data/ori_temp'


    if not os.path.exists(os.path.join(target_path,'fast_align_temp' )):
        os.mkdir(os.path.join(target_path,'fast_align_temp' ))
    fast_align_temp_path = os.path.join(target_path,'fast_align_temp' )

    pool = Pool(processes=30)
    files_img = os.listdir(img_path)
    pool.map(aling,files_img)

